"""Contains a variant of the densenet model definition."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf

slim = tf.contrib.slim


def trunc_normal(stddev): return tf.truncated_normal_initializer(stddev=stddev)

# Composite function 复合函数 BN->reLU->conv2d->dropout(0.5)
# num_outputs 输出数据channel数
# kernel_size 卷积核size
def bn_act_conv_drp(current, num_outputs, kernel_size, scope='block'):
    current = slim.batch_norm(current, scope=scope + '_bn')
    current = tf.nn.relu(current) # 此处的bn和ReLU实际是对上次卷积结果进行操作
    current = slim.conv2d(current, num_outputs, kernel_size, scope=scope + '_conv')
    current = slim.dropout(current, scope=scope + '_dropout')  # 默认值keep_prob=0.5
    return current

# Dense Block
# net 输入数据
# layers 层数
# growth 增长率
def block(net, layers, growth, scope='block'):
    for idx in range(layers):
        # 瓶颈层 Bottleneck layers
        bottleneck = bn_act_conv_drp(net, 4 * growth, [1, 1], scope=scope + '_conv1x1' + str(idx))
        tmp = bn_act_conv_drp(bottleneck, growth, [3, 3], scope=scope + '_conv3x3' + str(idx))
        # 每循环一层，拼接一层输出，即每过一层feature maps增加k个
        net = tf.concat(axis=3, values=[net, tmp]) # 在channel维度上进行拼接
    return net

# transition layer 变换层
def transition(net, num_outputs, scope='trans'):
    net = bn_act_conv_drp(net, num_outputs, [1,1], scope)
    net = slim.avg_pool2d(net, [2,2], stride=2, scope=scope + '_pooling')
    return net

# Classification Layer
def Classification(net, num_classes, scope='classification'):
    net = slim.avg_pool2d(net, kernel_size=net.shape[1:3], stride=1, scope=scope + 'global_avg_pool')
    net = slim.fully_connected(net, num_classes, scope=scope + 'Logits')
    return net

def densenet(images, num_classes=1001, is_training=False,
             dropout_keep_prob=0.8,
             scope='densenet'):
    """Creates a variant of the densenet model.

      images: A batch of `Tensors` of size [batch_size, height, width, channels].
      num_classes: the number of classes in the dataset.
      is_training: specifies whether or not we're currently training the model.
        This variable will determine the behaviour of the dropout layer.
      dropout_keep_prob: the percentage of activation values that are retained.
      prediction_fn: a function to get predictions out of logits.
      scope: Optional variable_scope.

    Returns:
      logits: the pre-softmax activations, a tensor of size
        [batch_size, `num_classes`]
      end_points: a dictionary from components of the network to the corresponding
        activation.
    """
    # Growth rate 增长率 k
    growth = 32
    # compression rate 压缩率 θ
    compression_rate = 0.5
    # block中层数
    layers = [6,12,24,16]
    # dense block数量
    num_dense_blocks = len(layers)

    # 计算要降的维数
    def reduce_dim(input_feature):
        return int(int(input_feature.shape[-1]) * compression_rate)

    end_points = {}

    with tf.variable_scope(scope, 'DenseNet', [images, num_classes]):
        with slim.arg_scope(bn_drp_scope(is_training=is_training,
                                         keep_prob=dropout_keep_prob)) as ssc:
            ##########################
            # 初始化，卷积+池化
            end_point = 'Initial'
            net = slim.conv2d(images, 2*growth, [7, 7], stride=2, scope=end_point + 'Conv2d_1_7x7')
            net = slim.batch_norm(net, scope=end_point + '_bn')
            net = tf.nn.relu(net)
            net = slim.max_pool2d(net, kernel_size=[3,3], stride=2, scope='MaxPool_1_3x3')
            end_points[end_point] = net

            for idx in range(num_dense_blocks):
                # Dense block
                end_point = 'Block_' + str(idx+1)
                net = block(net, layers[idx], growth, scope=end_point)
                end_points[end_point] = net

                if idx == (num_dense_blocks - 1):
                    break

                # transition layer
                end_point = 'Trans_' + str(idx+1)
                net = transition(net, reduce_dim(net), scope=end_point)
                end_points[end_point] = net

            # 此处需对最后一个block的卷积结果进行bn和ReLU的操作
            end_point = 'Last_block_BN_ReLU'
            net = slim.batch_norm(net, scope=end_point)
            net = tf.nn.relu(net)
            end_points[end_point] = net

            # Classification layer
            end_point = 'global_avg_pool'
            net = slim.avg_pool2d(net, kernel_size=net.shape[1:3], stride=1, scope=end_point)
            # net = slim.fully_connected(net, num_classes, scope=end_point + 'Logits')  # 为什么没有全连接？
            end_points[end_point] = net

            # softmax classifier
            end_point = 'Logits'
            logits = tf.squeeze(slim.conv2d(net, num_classes, [1, 1],
                                            biases_initializer=tf.constant_initializer(0.001),
                                            scope='logits'), [1, 2])
            end_points[end_point] = logits
            end_points['predictions'] = slim.softmax(logits, scope='predictions')
            ##########################

    return logits, end_points


def bn_drp_scope(is_training=True, keep_prob=0.8):
    keep_prob = keep_prob if is_training else 1
    with slim.arg_scope(
        [slim.batch_norm],
            scale=True, is_training=is_training, updates_collections=None):
        with slim.arg_scope(
            [slim.dropout],
                is_training=is_training, keep_prob=keep_prob) as bsc:
            return bsc


def densenet_arg_scope(weight_decay=0.04):
    """Defines the default densenet argument scope.

    Args:
      weight_decay: The weight decay to use for regularizing the model.

    Returns:
      An `arg_scope` to use for the inception v3 model.
    """
    with slim.arg_scope(
        [slim.conv2d],
        weights_initializer=tf.contrib.layers.variance_scaling_initializer(
            factor=2.0, mode='FAN_IN', uniform=False),
        activation_fn=None, biases_initializer=None, padding='same',
            stride=1) as sc:
        return sc


densenet.default_image_size = 224
